# encoding: utf-8


"""

@author: tongzhenguo

@time: 2021/5/16 下午2:21

@desc:

根据不同策略构建股票池(量级在几十个)
1.行业选股:宏观
2.分类选股:缓慢增长性、稳定增长型、快速增长型、周期型、困境反转型
3.特征选股:分拆上市

制定内容池条件
1.范围过滤(绝对值、同比增长率、比率)
市盈率倍数在(1,50)
2.全局黑名单：营业收入众多、主营业务占比小于60%、商誉占净资产比率>10%
3.ST不看
4.5年内新股不看
6.双创不看，买不了
7.可选，(月线)空头行情不买
"""
import sys

import numpy as np
import pandas as pd
from pandas import DataFrame

from app.stock_certificate_price import value_rectify, AppraisementCalculator
from app.technicals_analytics import MyTalib
from collect_stock_data import StockBase
from const import PROJ_HOME, TRADE_CAL_FILE, ALL_STOCK_A_FILE
from date_util import now_quarter, last_year, now_year, last_trade_date, now_date_fmt_8, get_n_day, last_n_year, \
    get_month_last_trade_date, gen_k_date_list, is_trade_date
from mongo_db import MongoDB


class StockSelection(StockBase):
    def __init__(self):
        super().__init__()

        self.stock_all_df = pd.read_csv(ALL_STOCK_A_FILE, encoding="utf-8", dtype={"code": str})
        self.all_stock_codes = list(self.stock_all_df["code"].values)

        black_list = set()
        self.mongo_db = MongoDB(uri="mongodb://127.0.0.1:27017", database="stock")
        # 年度财报时间
        self.year = int(last_year())
        # 季度财报时间
        self.quarter = int(now_quarter())
        # self.date = 20210622
        # 最近交易日时间，18点之后更新当日数据
        self.trade_cal_df = pd.read_csv(TRADE_CAL_FILE, dtype={"cal_date": str, "is_open": str})
        self.date = str(last_trade_date(self.trade_cal_df))
        df = self.mongo_db.fetch_data("stock", {}, ["_id", "%s日-滚动市盈率" % self.date, "商誉占净资产比例"])
        for col in ["%s日-滚动市盈率" % self.date, "商誉占净资产比例"]:
            df[col] = df[col].apply(np.float64)
        # NOTE 不包含市盈率为负的亏损公司
        values = df[df["%s日-滚动市盈率" % self.date] > 50.0]["_id"].values
        black_list.update(set(values))
        values = df[df["商誉占净资产比例"] > 0.1]["_id"].values
        black_list.update(set(values))
        # 创业板买不了
        values = df[df['_id'].str.startswith("3")]['_id'].values
        black_list.update(set(values))
        # 科创板买不了
        values = df[df['_id'].str.startswith("688")]['_id'].values
        black_list.update(set(values))

        # ST不看
        df = self.mongo_db.fetch_data("stock", {"isST": "1"}, ["_id"])
        values = df['_id'].values
        black_list.update(set(values))

        # 5年内新股不看
        n_year_5 = last_n_year(5, "str", format="YYYYMMDD")
        df = self.mongo_db.fetch_data("stock", {"上市日期": {"$gt": int(n_year_5)}}, ['_id'])
        values = df['_id'].values
        black_list.update(set(values))
        self.global_black_list = list(black_list)
        self.all_stock_codes = [c for c in self.all_stock_codes if c not in self.global_black_list]

        # 可选，(月线)空头行情不买
        start_date = "%s%s" % (last_n_year(2, mode="str", format="YYYY"), self.date[4:])
        end_date = get_month_last_trade_date(self.date[:4], self.date[4:6], self.trade_cal_df)
        columns = ["_id"]
        date_list = gen_k_date_list(start_date, end_date, self.trade_cal_df, frequency="m")
        columns.extend([f"{d}月-收盘价" for d in date_list])
        data = self.mongo_db.fetch_data("stock", {}, columns)
        for d in date_list:
            data[f"{d}月-收盘价"] = data[f"{d}月-收盘价"].apply(float)
        mt = MyTalib()
        data[f"{end_date}多头"] = data.apply(lambda r: mt.multi_head_order(r, date_list), axis=1)
        values = data[data[f"{end_date}多头"] == '']["_id"].values
        self.month_empty_black_list = list(values)

    def industry_select(self, industry, k=30, **kwargs):
        """获取行业topk股票池,
            总市值排名，总利润排名，总收入排名
        """
        df = self.mongo_db.fetch_data("stock", {"所属行业": industry},
                                      ["_id", "%s日-总市值" % self.date, "%s年度-利润总额" % self.year, "%s年度-营业总收入" % self.year])
        df = df[~df['_id'].isin(self.global_black_list)]
        if "month_empty_black_list" in kwargs and kwargs["month_empty_black_list"]:
            df = df[~df['_id'].isin(self.month_empty_black_list)]
        # 过滤掉利润和收入为负的
        for col in ["%s年度-利润总额" % self.year, "%s年度-营业总收入" % self.year]:
            df[col] = df[col].apply(np.float64)
            df = df[df[col] > 0]
        # 线性加权排名
        c_list = ["%s日-总市值" % self.date, "%s年度-利润总额" % self.year, "%s年度-营业总收入" % self.year]
        w_list = [0.5, 0.3, 0.1]
        for t in zip(c_list, w_list):
            df["score"] = t[1] * df[t[0]].apply(np.float64)
        df = df.sort_values("score", ascending=False).head(k)
        df.to_csv(PROJ_HOME + "/data/行业选股/%s.csv" % industry, index=False)
        print(df)
        return df

    def white_horse_stock_select(self, **kwargs):
        """蓝筹股和白马股
        白马股的特点：
        （1）每股收益在0.5元以上
        （2）每股净资产值3.00元以上
        （3）净资产收益率10%以上或者ROE连续7年大于等于15%
        （4）平均市盈率40倍以下,0以上
        （5）净利润增长率30%以上
        （6）主营业务收入增长率30%以上
        剔除股价过高的企业，过高可能会溢价。可以用市净率PB指标，30%的ROE对应的PB是8；20%对应的PB是4；
        负债率＜30%

        蓝筹股的特点:
        1. 利润率高
        2. 稳步成长
        3. 行业垄断
        """
        column_list = ["_id", "%sQ%s季度-摊薄每股收益(元)" % (now_year(), self.quarter),
                       "%sQ%s季度-每股净资产_调整后(元)" % (now_year(), self.quarter),
                       "%sQ%s季度-加权净资产收益率(%%)" % (now_year(), self.quarter),
                       "%sQ%s季度-净利润增长率(%%)" % (now_year(), self.quarter),
                       "%sQ%s季度-主营业务收入增长率(%%)" % (now_year(), self.quarter),
                       "2020年度-股票简称"]
        column_list.extend(["%s日-滚动市盈率" % self.date])
        df = self.mongo_db.fetch_data("stock", {}, column_list)
        df = df[~df['_id'].isin(self.global_black_list)]
        if "month_empty_black_list" in kwargs and kwargs["month_empty_black_list"]:
            df = df[~df['_id'].isin(self.month_empty_black_list)]
        column_list.remove("_id")
        # 筛选条件
        v_list = [0.5, 3, 10, 40, 30, 30]
        for t in zip(column_list, v_list):
            if t[0] == "2020年度-股票简称":
                continue
            col = t[0]
            if col == "%s日-滚动市盈率" % self.date:
                df[col] = df[col].apply(np.float64)
                df = df[df[col] <= t[1]]
                df = df[df[col] > 0]
                if df.shape[0] == 0:
                    print("%s数据可能异常" % col)
            else:
                df[col] = df[col].apply(np.float64)
                df = df[df[col] >= t[1]]
                if df.shape[0] == 0:
                    print("%s数据可能异常" % col)
        if df.shape[0] == 0:
            print("error jzx_select data empty !!!")
            return None
        df.to_csv(PROJ_HOME + "/data/价值白马/白马股.csv", index=False)
        print(df)
        return df

    def czx_select(self):
        """成长型选股，快速增长型;多半是中小盘
        """
        pass

    def kjfzx_select(self):
        """困境反转型"""
        pass

    def zqx_select(self):
        """周期型"""
        pass

    def ggx_select(self, k, **kwargs):
        """高股息率，红利
        各红利指数权重股,或股息率等指标前k
        todo 1.考虑使用连续5年的平均值做平滑；另外单纯看股息率
        todo 2.融资远大于分红的企业也要当心
        """
        # date + "日-滚动股息率"
        gxl_col = "%s日-滚动股息率" % self.date
        column_list = ["_id", gxl_col]
        df = self.mongo_db.fetch_data("stock", {}, column_list)
        df = df[~df['_id'].isin(self.global_black_list)]
        if "month_empty_black_list" in kwargs and kwargs["month_empty_black_list"]:
            df = df[~df['_id'].isin(self.month_empty_black_list)]
        # 如果股息率为空，重置为0
        df[gxl_col] = df[gxl_col].apply(lambda x: value_rectify(x, 0))
        df[gxl_col] = df[gxl_col].apply(np.float32)
        column_list.remove("_id")
        df = df.sort_values(gxl_col, ascending=False).head(k)
        if df.shape[0] == 0:
            print("error ggx_select data empty !!!")
            pass
        df.to_csv(PROJ_HOME + "/data/高股息红利/高股息红利.csv", index=False)
        print(df)
        return df

    def feature_select(self):
        """特征选股,比如分拆上市、股权激励、回购股票
        TODO 还没办法获取到这部分数据
        """
        pass

    def molenka_select(self, **kwargs):
        """莫伦卡选股
        指标1:过去5年来平均净资产收益率高于14%
        指标2:市盈率低于30 并且大于 0
        指标3:经营现金流为正
        指标4:新期的净利润大于前5年的净利润
        https://blog.csdn.net/sc_liuliye/article/details/113691256
        """
        column_list = ["_id"]
        for i in range(6):
            column_list.append("%s年度-净资产收益率(%%)" % (int(self.year) - i))
            column_list.append("%s年度-净利润" % (int(self.year) - i))
        column_list.append("%s日-滚动市盈率" % self.date)
        quarter_xjl_col = "%sQ%s季度-每股经营性现金流(元)" % (now_year(), self.quarter)
        column_list.append(quarter_xjl_col)
        # column_list.append("_id")
        df = self.mongo_db.fetch_data("stock", {}, column_list)
        df = df[~df['_id'].isin(self.global_black_list)]
        if "month_empty_black_list" in kwargs and kwargs["month_empty_black_list"]:
            df = df[~df['_id'].isin(self.month_empty_black_list)]
        df['过去5年来平均净资产收益率'] = 0
        for i in range(5):
            df['过去5年来平均净资产收益率'] = df['过去5年来平均净资产收益率'] + df["%s年度-净资产收益率(%%)" % (int(self.year) - i)].apply(np.float32)
        df['过去5年来平均净资产收益率'] = df['过去5年来平均净资产收益率'] / 5
        # 指标1
        df = df[df['过去5年来平均净资产收益率'] >= 14]
        if df.shape[0] == 0:
            print("过去5年来平均净资产收益率数据可能异常")
        # 指标2
        df["%s日-滚动市盈率" % self.date] = df["%s日-滚动市盈率" % self.date].apply(np.float64)
        print(df[["%s日-滚动市盈率" % self.date]])
        df = df[df["%s日-滚动市盈率" % self.date] <= 30]
        df = df[df["%s日-滚动市盈率" % self.date] > 0]
        if df.shape[0] == 0:
            print("%s日-滚动市盈率数据可能异常" % self.date)
        # 指标3
        df[quarter_xjl_col] = df[quarter_xjl_col].apply(np.float32)
        print(df[[quarter_xjl_col]])
        df = df[df[quarter_xjl_col] > 0]
        if df.shape[0] == 0:
            print("每股经营性现金流可能异常")
        # 指标4
        for i in range(1, 6):
            df['新期净利润-前5年的净利润'] = df["%s年度-净利润" % self.year].apply(np.float32) - df["%s年度-净利润" % (self.year - i)].apply(
                np.float32)
            df = df[df['新期净利润-前5年的净利润'] > 0]
            if df.shape[0] == 0:
                print("%s年度-净利润可能异常" % (self.year - i))
        if df.shape[0] == 0:
            print("error molenka_select data empty !!!")
            pass
        df.to_csv(PROJ_HOME + "/data/莫伦卡选股/莫伦卡选股.csv", index=False)
        print(df)
        return df

    def low_price(self, data: DataFrame, calculator, certificate_price_fn, val_col, lower_bound, date, epsion, k=30,
                  **kwargs):
        """低估值选股
        调用估值函数，并再考虑更低的安全边际，取价格低于价值epsion的前k支
        """
        func = getattr(calculator, certificate_price_fn)
        data = func(data)
        data[val_col] = data[val_col] * lower_bound
        data = data[data[val_col] != -1]
        if data.shape[0] == 0:
            return
        data['diff'] = data[val_col] - data['%s日-收盘价' % date]
        data = data[data['diff'] > epsion]
        df = data.sort_values('diff', ascending=False).head(k)
        if df.shape[0] == 0:
            print("error low_price data empty !!!")
            return
        df = df[~df['_id'].isin(self.global_black_list)]
        if "month_empty_black_list" in kwargs and kwargs["month_empty_black_list"]:
            df = df[~df['_id'].isin(self.month_empty_black_list)]
        df.to_csv(f"{PROJ_HOME}/data/低估值选股/{certificate_price_fn}.csv", index=False)
        print(df)
        return df

    def technology_select(self, **kwargs):
        """技术面选股
        1.形态选股(双重底)
        2.k线组合
        """

        def get_buy_signal(row):
            """优先形态买入信号，然后是k线组合"""
            if row["双重底(90自然日)"] in "双底形态突破有效":
                return "双底形态突破有效(90日)"
            elif row["双重底(90自然日)"] in "双底形态判断有效":
                return "双底形态判断有效(90日)"
            elif row["双重底(60自然日)"] == "双底形态突破有效":
                return "双底形态突破有效(60日)"
            elif row["双重底(60自然日)"] == "双底形态判断有效":
                return "双底形态判断有效(60日)"
            elif row["双重底(30自然日)"] == "双底形态突破有效":
                return "双底形态突破有效(30日)"
            elif row["双重底(30自然日)"] == "双底形态判断有效":
                return "双底形态判断有效(30日)"
            else:
                return row["买入信号"]

        column_list = ["2020年度-股票简称", "_id"]
        open_cols, high_cols, low_cols, close_cols = [], [], [], []
        for i in range(-90, 1):
            date = get_n_day(self.date, i)
            for p in ["开盘价", "最高价", "最低价", "收盘价"]:
                column_list.append(f"{date}日-{p}")
            open_cols.append(f"{date}日-开盘价")
            high_cols.append(f"{date}日-最高价")
            low_cols.append(f"{date}日-最低价")
            close_cols.append(f"{date}日-收盘价")
        df = self.mongo_db.fetch_data("stock", {}, column_list)
        column_list.remove("2020年度-股票简称")
        column_list.remove("_id")
        # for c in column_list:
        #     df[c] = df[c].apply(float)

        mt = MyTalib()
        df["双重底(30自然日)"] = df.apply(lambda x: mt.w_bottom(x, get_n_day(self.date, -30), self.date),
                                    axis=1)
        df["双重底(60自然日)"] = df.apply(lambda x: mt.w_bottom(x, get_n_day(self.date, -60), self.date),
                                    axis=1)
        df["双重底(90自然日)"] = df.apply(lambda x: mt.w_bottom(x, get_n_day(self.date, -90), self.date),
                                    axis=1)
        df = mt.history_k_pattern_buy_signal(df, open_cols, high_cols, low_cols, close_cols)
        df["买入信号"] = df.apply(get_buy_signal, axis=1)
        df = df[df["买入信号"] != "无"][["2020年度-股票简称", "_id", f"{self.date}日-收盘价", "买入信号"]]
        df = df[~df['_id'].isin(self.global_black_list)]
        if "month_empty_black_list" in kwargs and kwargs["month_empty_black_list"]:
            df = df[~df['_id'].isin(self.month_empty_black_list)]
        df.to_csv(PROJ_HOME + "/data/技术面选股/技术面选股.csv", index=False)

        print(df)


if __name__ == "__main__":
    ss = StockSelection()
    if len(sys.argv) > 1:
        date_s = sys.argv[1].strip()
        ss.date = date_s
    # 非交易日跳过
    if not is_trade_date(ss.date, ss.trade_cal_df):
        print(f"stock_selection 非交易日跳过")
        exit(0)

    calculator = AppraisementCalculator()
    data = ss.mongo_db.fetch_data("stock", {"_id": {"$in": ss.all_stock_codes}}, ["_id"])
    ss.low_price(data, calculator, "peterlynch_peg", "彼得林奇估值", 1.0 / 0.7, ss.date, 2, month_empty_black_list=True)
    ss.low_price(data, calculator, "pe_percent_rank", "pe分位点估值", 1., ss.date, 2, month_empty_black_list=True)

    ss.industry_select("新型电力", month_empty_black_list=True)
    ss.industry_select("软件服务", month_empty_black_list=True)
    ss.industry_select("通信设备", month_empty_black_list=True)
    ss.industry_select("银行")
    ss.industry_select("化学制药", month_empty_black_list=True)
    ss.industry_select("医疗保健", month_empty_black_list=True)
    ss.industry_select("互联网")
    ss.industry_select("半导体", month_empty_black_list=True)
    ss.industry_select("元器件", month_empty_black_list=True)

    ss.white_horse_stock_select()
    ss.ggx_select(30)
    ss.molenka_select()

    ss.technology_select(month_empty_black_list=True)

    # ss.feature_select()
    # ss.czx_select()
    # ss.kjfzx_select()
    # ss.zqx_select()
